#!/usr/bin/env python
# _*_ coding: utf-8 _*_
# @Project  :main.py
# @Date     :2021/1/23 11:12 下午
# @Author   :JIAN LI
# @Email    :lijianjlu@126.com

import cv2
import sys
import numpy as np
import datetime
import os
import glob
from retinaface import RetinaFace
import time
from retinaface_detection import retinaface_detect

if __name__ == '__main__':
    # 定义视频路径
    #video_path = 'xxxxx.avi' # 该路径可以是不同格式的视频，也可以是视频的链接地址
    video_path = 0 # 当路径设为0时，为打开本地摄像头

    # 初始化视频对象
    video = cv2.VideoCapture()

    if not video.open(video_path):
        print('视频加载失败！')

    # 加载RetinaFace模型
    gpuid = 0
    detector = RetinaFace('./models/retinaface-R50-cudnnoff/R50', 0, gpuid, 'net3')

    ret = True
    while ret:
        # 从视频中获取每帧图像
        ret, frame = video.read()

        # 针对每帧图像进行处理...
        # TO_DO(frame)

        # detect faces in the image
        faces, landmarks = retinaface_detect(frame, detector)

        # show the faces on the image
        if faces is not None:
            print('find', faces.shape[0], 'faces')
            for i in range(faces.shape[0]):
                #print('score', faces[i][4])
                box = faces[i].astype(np.int)
                #color = (255,0,0)
                color = (0,0,255)
                cv2.rectangle(frame, (box[0], box[1]), (box[2], box[3]), color, 2)
                if landmarks is not None:
                    landmark5 = landmarks[i].astype(np.int)
                    #print(landmark.shape)
                    for l in range(landmark5.shape[0]):
                        color = (0,0,255)
                        if l==0 or l==3:
                            color = (0,255,0)
                        cv2.circle(frame, (landmark5[l][0], landmark5[l][1]), 1, color, 2)

        # 显示每帧图像
        cv2.imshow('video',frame)

        # 定义退出视频方法，按"ESC"按键退出
        if cv2.waitKey(10) & 0xFF==27:
            break

    cv2.destroyAllWindows()

